import sensor, image, time, lcd, math
from machine import Timer, PWM,
from Maix import GPIO
from fpioa_manager import fm
# ============全局变量区==========
pins = {
    "LEDs": {
        "R": 30,
        "G": 31,
        "Y": 32
    },
    "LOADED": 33,
    "MOTORs": {
        "1_1": 20,
        "1_2": 21,
        "2_1": 22,
        "2_2": 23
    }
}
timer_motor = Timer.TIMER0
channel_motor1_1 = Timer.CHANNEL0
channel_motor1_2 = Timer.CHANNEL1
channel_motor2_1 = Timer.CHANNEL2
channel_motor2_2 = Timer.CHANNEL3

# ==============数字检测类=======
# 使用方法：
# detector = DiditDetector()
# nums = detector.detect(img)
class DigitDetector():
    def __init__(self):
        self.thresholds = [(0, 40, -128, 127, -128, 127)]#LAB阈值
        self.templates = []
        for i in range(1, 9):
            self.templates.append(image.Image("/sd/%d.bmp" % i).to_grayscale())
    def detect(self, img, draw_rect=True):
        nums = []
        # 先寻找可能是数字的黑色色块
        for b in img.find_blobs(self.thresholds, pixels_threshold=50, area_threshold=160, merge=False):
            # 通过形状先排除掉一部分色块
            if b.h() / b.w() > 1 and b.h() / b.w() <= 4:
                # 循环匹配每一个模板
                for idx, template in enumerate(self.templates):
                    # 如果之前就匹配到了这个数字了，就跳过
                    if (idx+1) in nums:
                        continue
                    # 重设模板尺寸，使其等于候选区尺寸
                    template = template.resize(b.w(), b.h())
                    # 模板匹配
                    res = img.to_grayscale().find_template(
                                                template, 0.8,
                                                step=1, search=image.SEARCH_EX,
                                                roi=(b.x()-2, b.y()-2, b.w()+4, b.h()+4))
                    # 如果有结果，那就添加到nums，然后结束这一轮匹配
                    if res:
                        # 在img表面绘制方框和数字
                        if draw_rect:
                            img.draw_rectangle(b.rect())
                            img.draw_string(b.x()+5, b.y()+5, str(idx+1))
                        nums.append((idx+1, b.cx()))
                        break
        # 最后返回匹配结果
        sorted(nums, key=lambda x: x[1])
        reversed(nums)
        nums = list(map(lambda x: x[0], nums))
        return nums
# =============数字检测类结束=============
# ==============巡线类====================
class LineDetector():
    def __init__(self):
        # 巡线ROI
        self.roi_line = [
            (5, 15, 150, 20),  # 上
            (5, 48, 150, 20), # 中
            (5, 86, 150, 20)  # 下
        ]
        # 检测十字或T字ROI
        self.roi_cross = [
            (0, 0, 32, 60),  # 左
            (128, 0, 32, 60) # 右
        ]
        # 检测终点的方块的ROI
        self.roi_end = (0, 30, 160, 60)
        self.th = (0, 35, 5, 48, -11, 33)
    def check(self, img):
        # 对竖直线进行检测
        pos_v = []
        for roi in self.roi_line:
            blobs = img.find_blobs([self.th], roi=roi, area_threshold=16)
            avg = 0
            for b in blobs:
                avg += b.cx()
                img.draw_cross(b.cx(), b.cy(), size=5)
            if len(blobs) != 0:
                avg /= len(blobs)
                pos_v.append(avg/160)
            if len(pos_v) == 3:
                rotation = pos_v[0] - pos_v[2]  # 正右负左
                motion = pos_v[1] - 0.5  # 如果是正的则表示向右偏了，负的就是往左偏了
                img.draw_arrow(80, 60, 80 + int(motion*200), 60)
                img.draw_line(int(pos_v[0]*160), 12, int(pos_v[2]*160), 108)
        # 对水平线进行检测
        pos_h = []
        # 遍历两个ROI区域
        for roi in self.roi_cross:
            img.draw_rectangle(roi[0], roi[1], roi[2], roi[3])
            # 在ROI内寻找轨迹的颜色
            blobs = img.find_blobs([self.th], roi=roi, area_threshold=16)
            # ROI内所有符合要求的色块的中心纵坐标的平均值
            avg = 0
            # 遍历检测到的所有色块，计算纵坐标平均值
            for b in blobs:
                avg += b.cy()
            if len(blobs) != 0:
                avg /= len(blobs)
                pos_h.append(avg/120)
            # 如果左右都检测到了轨道，那么把水平线画出来
            if len(pos_h) == 2:
                img.draw_cross(8, int(pos_h[0]*120), size=5)
                img.draw_cross(152, int(pos_h[1]*120), size=5)
                img.draw_line(8, int(pos_h[0]*120), 152, int(pos_h[1]*120))
        # 检测终点
        rects = img.find_rects(self.roi_end, 5000)
        for rect in rects:
            img.draw_rectangle(rect.x(), rect.y(), rect.w(), rect.h())
        # 当前的状态变量
        status = {}
        # 当前小车没有遇到交叉点
        if len(pos_v) == 3 and len(pos_h) < 2 and len(rects) <= 5:
            status["pos"] = "line"
            status["x"] = pos_v[1]              # 当前偏离轨道多少
            status["dx"] = pos_v[0]-pos_v[2]    # 当前旋转了多少
        # 只要检测水平线的ROI检测到了两个数据，那一定是在交叉点
        elif len(pos_h) == 2  and len(rects) <= 5:
            status["pos"] = "cross"
            status["y"] = (pos_h[0]+pos_h[1])/2
        elif len(rects) >= 5:
            status["pos"] = "end"
        return status
# ========巡线类结束=========
# =========PID控制类=========
class PIDControler:
    def __init__(self, p, i, d):
        self.kp = p  # 比例系数
        self.ki = i  # 积分系数
        self.kd = d  # 微分系数
        self.err = 0
        self.prev_err = 0
        self.p = 0
        self.i = 0
        self.d = 0
    def pid(self, err, dt):
        self.p = self.kp * err
        self.i += self.ki * err * dt
        self.d = self.kd * ((err - self.prev_err)/dt)
        self.prev_err = err
        return self.p + self.i - self.d
    def zero(self):
        self.err = 0
        self.prev_err = 0
        self.p = 0
        self.i = 0
        self.d = 0
# =====PID控制类结束=======
# ====运动控制类=========
class MotionManager():
    def __init__(self):
        # 初始化控制电机用的4路PWM时钟
        self.timer_pwm_1_1 = Timer(timer_motor, channel_motor1_1, mode = Timer.MODE_PWM)
        self.timer_pwm_1_2 = Timer(timer_motor, channel_motor1_2, mode = Timer.MODE_PWM)
        self.timer_pwm_2_1 = Timer(timer_motor, channel_motor2_1, mode = Timer.MODE_PWM)
        self.timer_pwm_2_2 = Timer(timer_motor, channel_motor2_2, mode = Timer.MODE_PWM)
        # 初始化4路PWM
        self.pwm11 = PWM(self.timer_pwm_1_1, 1000, 0, pins["MOTORs"]["1_1"])
        self.pwm12 = PWM(self.timer_pwm_1_2, 1000, 0, pins["MOTORs"]["1_2"])
        self.pwm21 = PWM(self.timer_pwm_2_1, 1000, 0, pins["MOTORs"]["2_1"])
        self.pwm22 = PWM(self.timer_pwm_2_2, 1000, 0, pins["MOTORs"]["2_2"])
        # 初始化数字识别
        self.digit_detector = DigitDetector()
        # 初始化巡线
        self.line_detector = LineDetector()
        # 初始化PID控制器
        self.pid = PIDControler(1, 0.1, 0.1)
        # 注册引脚
        fm.register(pins["LEDs"]["R"], fm.fpioa.GPIO0)
        fm.register(pins["LEDs"]["G"], fm.fpioa.GPIO1)
        fm.register(pins["LEDs"]["Y"], fm.fpioa.GPIO2)
        fm.register(pins["LOADED"], fm.fpioa.GPIO3)

        self.led_r = GPIO(GPIO.GPIO0, GPIO.OUT)
        self.led_g = GPIO(GPIO.GPIO1, GPIO.OUT)
        self.led_y = GPIO(GPIO.GPIO2, GPIO.OUT)
        self.pin_load = GPIO(GPIO.GPIO3, GPIO.IN, GPIO.PULL_DOWN)
    # 等待药物
    def wait_medicine(self, state):
        print("等待药物状态: %d" % state)
        while(self.pin_load.value() != state):
            pass
    # 控制灯
    def light_state(self, led, state):
        leds = {
            "R": self.led_r,
            "G": self.led_g,
            "Y": self.led_y
        }
        leds[led].value(state)
    # 停止所有PWM通道的输出
    def stop(self):
        self.pwm11.duty(0)
        self.pwm12.duty(0)
        self.pwm21.duty(0)
        self.pwm22.duty(0)
    def motor_stop(self, motor):
        if motor == 1:
            self.pwm11.duty(0)
            self.pwm12.duty(0)
        else:
            self.pwm21.duty(0)
            self.pwm22.duty(0)
    # 控制某个电机按照特定速度运行
    def motor_go(self, motor, speed):
        self.stop()
        # 电机1
        if motor == 1:
            if speed > 0:
                self.pwm11.duty(speed)
            elif speed < 0:
                self.pwm12.duty(-speed)
            else:
                self.pwm11.duty(0)
                self.pwm12.duty(0)
        # 电机2
        else:
            if speed > 0:
                self.pwm21.duty(speed)
            elif speed < 0:
                self.pwm22.duty(-speed)
            else:
                self.pwm21.duty(0)
                self.pwm22.duty(0)
    # 开始前进
    def forward(self):
        print("开始前进")
        self.motor_go(1, 1)
        self.motor_go(2, 1)
    # 移动到交叉点
    def move_to_cross(self):
        img = sensor.snapshot()
        state = self.line_detector.check(img)
        # 开始前进
        self.forward()
        # 等待检测不到交叉点
        while ("pos" in state) and (state["pos"] == "cross"):
            img = sensor.snapshot()
            state = self.line_detector.check(img)
        print("已离开交叉点")
        time.sleep_ms(500)
        # 等待检测到交叉点
        while ("pos" not in state) or state["pos"] != "cross" or state["pos"] == "cross" and state["y"] < 0.15:
            img = sensor.snapshot()
            state = self.line_detector.check(img)
        # print(state)
        # 停止
        self.stop()
        print("已移动到交叉点")
    # 转向
    def turn(self, direction):
        delta_time = 500
        # 左转
        if direction == 0:
            self.motor_go(1, 1)
            self.motor_go(2, -1)
            print("左转")
        # 右转
        elif direction == 1:
            self.motor_go(1, -1)
            self.motor_go(2, 1)
            print("右转")
        # 向后转
        else:  # == 2
            self.motor_go(1, 1)
            self.motor_go(2, -1)
            print("向后转")
        # 延时
        if direction != 2:
            time.sleep_ms(delta_time)
        else:
            time.sleep_ms(2*delta_time)
    # 移动到终点
    def move_to_end(self):
        state = {}
        # 等待检测到交叉点
        while ("pos" not in state) or state["pos"] != "end":
            img = sensor.snapshot()
            state = self.line_detector.check(img)
        # 停止
        self.stop()
        print("移动到了终点")
    # 检测指定数量的数字
    def detect_nums(self, num_nums):
        print("开始识别数字")
        nums = []
        while len(nums) != num_nums:
            img = sensor.snapshot()
            nums = self.digit_detector.detect(img)
        return nums
    # 前往非近端病房
    def goto_room(self, aim, stack, depth):
        self.move_to_cross()
        print("移动到cross%d" % depth)
        stack.append("GO")
        if depth == 1:  # 第一个交叉点，没有数字，移动到下一个
            self.goto_room(aim, stack, depth + 1)
            return
        elif depth == 2:  # 第二个交点，有两个数字
            nums = self.detect_nums(2)
            print(nums)
            # 判断目标数字在不在识别到的数字里
            if aim in nums:
                # 在，接下来判断数字在那边
                pos = nums.index(aim)
                if pos == 0:  #  在左边
                    self.turn(0)
                    stack.append("TL")
                else:  # 在右边
                    self.turn(1)
                    stack.append("TR")
                # 移动到终点
                self.move_to_end()
                stack.append("GO")
                return
            else:  # 不在，那就递归这个函数
                self.goto_room(aim, stack, depth + 1)
                return
        elif depth == 3:  # 第三个交点，有4个数字
            nums = self.detect_nums(4)
            print(nums)
            # 目标一定在四个数字里
            pos = nums.index(aim)
            if pos == 1 or pos == 2:  # 数字在左边
                self.turn(0)
                stack.append("TL")
            else:
                self.turn(1)
                stack.append("TR")
            self.goto_room(aim, stack, depth + 1)
            return
        else:  # 第四个交点，有2个数字
            nums = self.detect_nums(2)
            print(nums)
            pos = nums.index(aim)
            if pos == 0:  # 在左边
                self.turn(0)
                stack.append("TL")
            else:
                self.turn(1)
                stack.append("TR")
            # 移动到终点
            self.move_to_end()
            stack.append("GO")
            return
    # 控制小车返回
    def backward(self, forward):
        reversed(forward)
        # 调头
        self.turn(2)
        # 反向寻径
        for action in forward[:-1]:
            if action == "GO":
                self.move_to_cross()
            elif action == "TL":
                self.turn(1)
            elif action == "TR":
                self.turn(0)
        self.move_to_end()
    # 开始运行
    def run(self):
        # 保存路径
        stack = []
        # 1.检测模板房间数字
        num_aim = self.detect_nums(1)[0]
        print("目标房间：%d" % num_aim)
        self.wait_medicine(1)
        # 2.判断是否为近端病房
        # 近端病房
        if num_aim == 1 or num_aim == 2:
            print("近端")
            # 到交叉点
            self.move_to_cross()
            stack.append("GO")
            if num_aim == 1:
                self.turn(0)  # 左转
                stack.append("TL")
            else:
                self.turn(1)  # 右转
                stack.append("TR")
            self.move_to_end()  # 到终点
            stack.append("GO")
        # 非近端病房
        else:
            self.goto_room(num_aim, stack, 1)
        print(stack)
        self.wait_medicine(0)
        self.backward(stack)
# ==运动控制类结束==
def system_init():
    sensor.reset()
    sensor.set_pixformat(sensor.RGB565)
    sensor.set_framesize(sensor.QQVGA)
    sensor.skip_frames(30)
    sensor.set_auto_gain(False)
    sensor.set_auto_whitebal(False)
    lcd.init()

if __name__ == "__main__":
    system_init()
    mm = MotionManager()
    mm.run()

